#!/usr/bin/env python3
import asyncio
import json
import os
import signal
import sys
import websockets


users = {}

def signal_handler(signal, frame):
    sys.exit(0)


async def error(websocket, message):
    """
    Send an error message.

    """
    event = {
        "type": "error",
        "message": message,
    }
    await websocket.send(json.dumps(event))


async def play(websocket):
    """
    收发消息
    """
    async for message in websocket:
        event = json.loads(message)
        try:
            if 'heart' in event['type']:
                event = {
                    "type":"heart",
                    "data":"ok"
                    }
                #await websocket.send(json.dumps(event))
                pass
            else:
                 await websocket.send(json.dumps(event))
        except RuntimeError as exc:
            await error(websocket, str(exc))
            continue


async def start(websocket,data):
    """
    认证检查api_token是否合法不合法关闭链接

    """
    api = data["api_token"]
    users[api] = {'websocket':websocket}
    print(users)
    try:
        event = {
            "type": "init",
            "data" :"ok"
        }
        #await websocket.send(json.dumps(event))
        #连接后发送消息
        await play(websocket)
    except websockets.exceptions.ConnectionClosedError as e:
        pass
    finally:
        del users[api]
        #print("del token：" + api)


async def publish(websocket,data):
    """
    推送消息给手机

    """
    try:
        event = {
            "type": "publish",
            "data" :"ok"
        }
        print(data)
        api_token = data['api_token']
        if api_token in users:
            await users[api_token]['websocket'].send(json.dumps(data["data"]))
    finally:
       pass




async def handler(websocket):
    """
    处理认证和推送

    """
    message = await websocket.recv()
    event = json.loads(message)
    if "publish" in event['type']:
        await publish(websocket, event)
    else:
        await start(websocket,event)


async def main():
    # 主服务
    loop = asyncio.get_running_loop()
    stop = loop.create_future()
    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)

    port = int(os.environ.get("PORT", "8765"))
    async with websockets.serve(handler, "0.0.0.0", port):
        await stop
        


if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal_handler)
    asyncio.run(main())